from ycrx import 计算退休时间, 计算法定退休时间_sim  # 假设这是正确的导入路径
import pandas as pd
from datetime import datetime
from openpyxl import Workbook
from openpyxl.styles import PatternFill, Border, Side, Alignment

if __name__ == "__main__":
    测试数据 = [
        {"分类": "男", "开始时间": (1965 - 5, 1), "结束时间": (1977 + 5, 1)},
        {"分类": "女干部", "开始时间": (1970 - 5, 1), "结束时间": (1982 + 5, 1)},
        {"分类": "女职工", "开始时间": (1975 - 5, 1), "结束时间": (1985 + 5, 1)}
    ]

    时间戳 = datetime.now().strftime("%Y%m%d%H%M%S")
    文件名 = f"退休时间速查表_{时间戳}.xlsx"

    wb = Workbook()

    for 数据 in 测试数据:
        分类 = 数据["分类"]
        开始时间 = 数据["开始时间"]
        结束时间 = 数据["结束时间"]
        数据列表 = []
        行号 = 1

        print(f"测试 {分类} 类别:")
        当前年, 当前月 = 开始时间
        while (当前年, 当前月) < 结束时间:
            结果 = 计算退休时间(分类, 当前年, 当前月)
            结果sim = 计算法定退休时间_sim(分类, 当前年 * 100 + 当前月)  # 调用计算函数获取结果2
            退休时年龄 = f"{结果['退休时年龄'][0]}岁"
            if 结果['退休时年龄'][1] != 0:
                退休时年龄 += f"{结果['退休时年龄'][1]}个月"
            延迟时长 = "不涉及"
            if 结果['延迟退休时长'] is not None:
                延迟时长 = ""
                if 结果['延迟退休时长'][0] != 0:
                    延迟时长 += f"{结果['延迟退休时长'][0]}年"
                if 结果['延迟退休时长'][1] != 0:
                    延迟时长 += f"{结果['延迟退休时长'][1]}个月"
            弹性提前 = "不涉及"  # 初值设为"不涉及"
            if 结果['弹性提前退休限制'] is not None:
                弹性提前 = "%04d-%02d" % 结果['弹性提前退休限制']

            数据列表.append([
                行号,
                "%04d-%02d" % (当前年, 当前月),
                退休时年龄,
                "%04d-%02d" % 结果["法定退休时间"],
                延迟时长,
                弹性提前,
                "%04d-%02d"%(结果sim//100,结果sim % 100)  # 添加这一行，将结果2添加到数据列表中
            ])

            print("%03d" % 行号, 分类, "出生:%04d-%02d" % (当前年, 当前月), f"退休时年龄:{退休时年龄}",
                  "法定退休时间:%04d-%02d" % 结果["法定退休时间"], f"延迟:{延迟时长}", f"弹性提前退休:{弹性提前}")
            行号 += 1
            当前月 += 1
            if 当前月 > 12:
                当前月 = 1
                当前年 += 1
        print("-" * 40)

        df = pd.DataFrame(数据列表, columns=["行号", "出生年月", "退休时年龄", "法定退休年月", "延迟", "弹性提前退休", "法定退休年月_sim"])

        ws = wb.create_sheet(title=分类)

        # 插入列标题
        ws.append(df.columns.tolist())

        # 添加DataFrame中的数据
        for r in df.iterrows():
            ws.append(r[1].tolist())

        # 设置列宽
        ws.column_dimensions['A'].width = 5
        ws.column_dimensions['B'].width = 16  # 出生年月
        ws.column_dimensions['C'].width = 12  # 退休时年龄
        ws.column_dimensions['D'].width = 15  # 法定退休年月
        ws.column_dimensions['E'].width = 12
        ws.column_dimensions['F'].width = 12
        ws.column_dimensions['G'].width = 15  # 新增的“法定退休年月_sim”列
        ws.column_dimensions['H'].width = 10  # 新增的“比较结果”列

        # 应用表格线
        thin_border = Border(left=Side(style='thin'),
                             right=Side(style='thin'),
                             top=Side(style='thin'),
                             bottom=Side(style='thin'))

        for row in ws.iter_rows():
            for cell in row:
                cell.border = thin_border

        # 为H列添加公式
        for i in range(2, ws.max_row + 1):
            col_d = f"D{i}"
            col_g = f"G{i}"
            formula = f'=IF({col_d}={col_g}, "ok", "err")'
            ws[f"H{i}"] = formula

        # 背景颜色处理
        light_yellow_fill = PatternFill(start_color="FFFFCC", end_color="FFFFCC", fill_type="solid")
        grey_fill = PatternFill(start_color="D3D3D3", end_color="D3D3D3", fill_type="solid")

        for row in ws.iter_rows(min_row=2, max_col=ws.max_column, max_row=ws.max_row):
            if row[4].value != "不涉及":
                for cell in row:
                    cell.fill = light_yellow_fill
            elif row[4].value == "不涉及":
                for cell in row:
                    cell.fill = grey_fill

    # 删除默认创建的工作表
    del wb[wb.sheetnames[0]]

    # 保存文件
    wb.save(文件名)
    print(f"已生成文件：{文件名}")